import math

import numpy as np
import matplotlib.pyplot as plt

def de_casteljau(t, ctrl_pts):
    """
    de Casteljau 递归实现
    参数
    ----
    t : float  (0 ≤ t ≤ 1)
        曲线参数
    ctrl_pts : ndarray, shape (n+1, 2)
        控制点列表，n 为 Bézier 曲线次数
    返回
    ----
    point : ndarray, shape (2,)
        参数 t 对应的曲线点
    """
    pts = np.array(ctrl_pts, dtype=float)
    n = pts.shape[0] - 1          # 曲线次数
    # 逐层线性插值
    for r in range(1, n + 1):
        pts[:n - r + 1] = (1 - t) * pts[:n - r + 1] + t * pts[1:n - r + 2]
    return pts[0]

def rational_de_casteljau(P, w, t):
    # P: list of 2D control points [(x0,y0),..., (x3,y3)]
    # w: list of weights [w0,w1,w2,w3]
    # 构造齐次坐标
    Q = [(w[i]*P[i][0], w[i]*P[i][1], w[i]) for i in range(4)]

    # 递归插值
    for r in range(1, 4):               # 3 层
        Q = [( (1-t)*Q[i][0] + t*Q[i+1][0],
               (1-t)*Q[i][1] + t*Q[i+1][1],
               (1-t)*Q[i][2] + t*Q[i+1][2]) for i in range(4-r)]

    X, Y, W = Q[0]
    return (X/W, Y/W)

if __name__ == '__main__':
    # 端点和控制点
    P0 = np.array([1.0, 0.0])
    P1 = np.array([2.0, -1.0])
    P2 = np.array([3.0, -1.0])
    P3 = np.array([4.0, 2.0])
    w = [2, 1, 1, 2]

    # 生成 t 参数
    t_vals = np.linspace(0, 1, 200)
    ctrl_pts = [P0, P1, P2, P3]

    # 绘图
    plt.figure(figsize=(6, 6))

    # 贝塞尔曲线
    curve = np.array([de_casteljau(t, ctrl_pts) for t in t_vals])
    plt.plot(curve[:, 0], curve[:, 1], 'b-', label='Bezier curve')

    # 有理贝塞尔曲线
    rational_curve = np.array([rational_de_casteljau(ctrl_pts, w, t) for t in t_vals])
    plt.plot(rational_curve[:,0], rational_curve[:,1], 'r-', label='Rational Bezier curve')

    # 控制多边形
    control_pts = np.vstack(ctrl_pts)
    plt.plot(control_pts[:, 0], control_pts[:, 1], 'k--', marker='o',
             label='Control polygon')
    plt.title('Cubic Bézier Curve')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.axis('equal')
    plt.grid(True)
    plt.legend()
    plt.show()
